import os
import shutil
import random
import torch
import cv2
from tqdm import tqdm
from ultralytics import YOLO

# 加载YOLOv8的模型
model = YOLO('yolo11n.pt')

source = r'Z:\Code_Pile\Programs\DataProcess\Data\data_origin'
target_path = r'Z:\Code_Pile\Programs\DataProcess\Data\2025_data_cleaned\origin'
if not os.path.exists(target_path):
    os.makedirs(target_path)
else:
    print(target_path)

# 遍历视频文件
for video in tqdm([f for f in os.listdir(source) if f.endswith(('.mp4', '.avi', '.mov'))]):
    video_path = os.path.join(source, video)
    video_capture = cv2.VideoCapture(video_path)
    total_frames = int(video_capture.get(cv2.CAP_PROP_FRAME_COUNT))

    # 检查帧数是否足够
    if total_frames < 1:
        video_capture.release()
        continue

    random_frames = random.sample(range(total_frames), min(3, total_frames))  # 最多检查3帧
    check = [True] * len(random_frames)
    having_ship = False

    for i, frame_number in enumerate(random_frames):
        video_capture.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
        success, frame = video_capture.read()
        if not success:
            check[i] = False
            break

        results = model.predict(frame, conf=0.5, classes=8)
        for r in results:
            if len(r.boxes.cls) != 1:  # 检查类别数
                check[i] = False
                break
            if r.boxes.cls.item() == 8:  # 检查类别是否为8
                having_ship = True

    video_capture.release()

    # 如果检测符合条件，复制文件
    if all(check) and having_ship:
        shutil.copy(video_path, os.path.join(target_path, video))

